#!/usr/bin/env python3
"""
LaTeX Error Fixer Script
Fixes common LaTeX compilation errors in batch for the C++ Function Call Tree report.
"""

import os
import re
import glob
from pathlib import Path

def fix_algorithm_commands(content):
    """Fix algorithmicx command issues"""
    # Fix ElsIf to ElsIf (correct algorithmicx command)
    content = re.sub(r'\\ElsIf\{', r'\\ElsIf{', content)
    
    # Fix common algorithm command issues
    content = re.sub(r'\\Elsif\{', r'\\ElsIf{', content)
    content = re.sub(r'\\ElseIf\{', r'\\ElsIf{', content)
    
    # Ensure proper If/EndIf matching
    content = re.sub(r'\\EndIf(?!\s*\\)', r'\\EndIf\n', content)
    content = re.sub(r'\\EndFor(?!\s*\\)', r'\\EndFor\n', content)
    content = re.sub(r'\\EndWhile(?!\s*\\)', r'\\EndWhile\n', content)
    
    return content

def fix_floating_environments(content):
    """Fix 'Not in outer par mode' errors by adding [H] to floating environments"""
    # Fix algorithm environments
    content = re.sub(r'\\begin\{algorithm\}(\s*)', r'\\begin{algorithm}[H]\\1', content)
    
    # Fix table environments in problematic contexts
    content = re.sub(r'\\begin\{table\}(\s*)', r'\\begin{table}[H]\\1', content)
    
    # Fix figure environments in problematic contexts
    content = re.sub(r'\\begin\{figure\}(\s*)', r'\\begin{figure}[H]\\1', content)
    
    return content

def fix_math_mode_issues(content):
    """Fix mathematical expressions and brace issues"""
    # Fix common math mode issues with Chinese text
    content = re.sub(r'\$([^$]*[\u4e00-\u9fff][^$]*)\$', lambda m: '$' + m.group(1).replace(' ', r'\ ') + '$', content)
    
    # Fix missing braces in mathematical expressions
    content = re.sub(r'(\\\w+)\s*\{([^}]*)\s*$', r'\1{\2}', content, flags=re.MULTILINE)
    
    return content

def fix_caption_issues(content):
    """Fix caption placement issues"""
    # Ensure captions come after begin{algorithm} and before begin{algorithmic}
    pattern = r'(\\begin\{algorithm\}(?:\[H\])?)\s*(\\caption\{[^}]*\})\s*(\\label\{[^}]*\})?\s*(\\begin\{algorithmic\})'
    replacement = r'\1\n\2\n\3\n\4'
    content = re.sub(pattern, replacement, content)
    
    return content

def fix_brace_matching(content):
    """Attempt to fix basic brace matching issues"""
    lines = content.split('\n')
    fixed_lines = []
    
    for line in lines:
        # Count braces in algorithmic environments
        if '\\begin{algorithmic}' in line or '\\end{algorithmic}' in line:
            fixed_lines.append(line)
            continue
            
        # Fix obvious brace mismatches in algorithm lines
        if any(cmd in line for cmd in ['\\State', '\\If', '\\For', '\\While', '\\Return', '\\Comment']):
            # Count opening and closing braces
            open_braces = line.count('{')
            close_braces = line.count('}')
            
            # Simple heuristic: if we have unclosed braces at end of line, close them
            if open_braces > close_braces and line.strip().endswith('}') == False:
                # Don't add braces if line ends with a comment or specific constructs
                if not re.search(r'\\Comment\{[^}]*$', line) and not line.strip().endswith('\\'):
                    line += '}'
        
        fixed_lines.append(line)
    
    return '\n'.join(fixed_lines)

def fix_chinese_text_issues(content):
    """Fix issues with Chinese text in LaTeX"""
    # Ensure proper spacing around Chinese text in mathematical contexts
    content = re.sub(r'(\$[^$]*)([\u4e00-\u9fff]+)([^$]*\$)', r'\1\\text{\2}\3', content)
    
    # Fix Chinese text in comments
    content = re.sub(r'\\Comment\{([^}]*[\u4e00-\u9fff][^}]*)\}', r'\\Comment{\1}', content)
    
    return content

def fix_table_environment_conflicts(content):
    """Fix table environment conflicts"""
    # Ensure tables are not nested in inappropriate environments
    # Add proper spacing around tables
    content = re.sub(r'(\\end\{[^}]+\})\s*(\\begin\{table\})', r'\1\n\n\2', content)
    content = re.sub(r'(\\end\{table\})\s*(\\begin\{[^}]+\})', r'\1\n\n\2', content)
    
    return content

def process_file(filepath):
    """Process a single LaTeX file"""
    print(f"Processing {filepath}")
    
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            content = f.read()
        
        original_content = content
        
        # Apply fixes in order
        content = fix_algorithm_commands(content)
        content = fix_floating_environments(content)
        content = fix_caption_issues(content)
        content = fix_math_mode_issues(content)
        content = fix_chinese_text_issues(content)
        content = fix_brace_matching(content)
        content = fix_table_environment_conflicts(content)
        
        # Only write if content changed
        if content != original_content:
            # Create backup
            backup_path = str(filepath) + '.backup'
            with open(backup_path, 'w', encoding='utf-8') as f:
                f.write(original_content)
            print(f"  Created backup: {backup_path}")
            
            # Write fixed content
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(content)
            print(f"  Fixed issues in {filepath}")
            return True
        else:
            print(f"  No changes needed for {filepath}")
            return False
            
    except Exception as e:
        print(f"  Error processing {filepath}: {e}")
        return False

def main():
    """Main function"""
    print("LaTeX Error Fixer Script")
    print("=" * 50)
    
    # Get current directory
    base_dir = Path('.')
    
    # Find all .tex files
    tex_files = []
    
    # Main document
    main_doc = base_dir / 'cpp-function-call-tree-report.tex'
    if main_doc.exists():
        tex_files.append(main_doc)
    
    # Chapter files
    chapters_dir = base_dir / 'chapters'
    if chapters_dir.exists():
        tex_files.extend(chapters_dir.glob('*.tex'))
    
    if not tex_files:
        print("No .tex files found!")
        return
    
    print(f"Found {len(tex_files)} .tex files to process:")
    for f in tex_files:
        print(f"  {f}")
    
    print("\nProcessing files...")
    print("-" * 30)
    
    processed_count = 0
    fixed_count = 0
    
    for tex_file in tex_files:
        processed_count += 1
        if process_file(tex_file):
            fixed_count += 1
    
    print("-" * 30)
    print(f"Processing complete!")
    print(f"Files processed: {processed_count}")
    print(f"Files with fixes: {fixed_count}")
    
    # Try to compile after fixes
    print("\nAttempting compilation...")
    print("-" * 30)
    
    if main_doc.exists():
        import subprocess
        try:
            result = subprocess.run(
                ['xelatex', '-interaction=nonstopmode', main_doc.name],
                cwd=main_doc.parent,
                capture_output=True,
                text=True,
                timeout=60
            )
            
            if result.returncode == 0:
                print("✅ Compilation successful!")
            else:
                print("⚠️  Compilation completed with errors/warnings")
                print("Check the .log file for details")
                
        except subprocess.TimeoutExpired:
            print("⏱️  Compilation timed out")
        except FileNotFoundError:
            print("❌ xelatex not found. Please install XeLaTeX")
        except Exception as e:
            print(f"❌ Compilation error: {e}")
    
    print("\nDone! Check the output and .backup files if you need to restore.")

if __name__ == '__main__':
    main()